'''
datasets_map = {
    'dc':{
        'ori_dir':'/mnt/nas/shengjie/datasets/DressCode_1024',
        'types':[types[0],types[1],types[2],],
        'clo_dir':['upper/cloth','lower/cloth','dresses/cloth'],
        'human_dir':['upper/image','lower/image','dresses/image'],
        'get_human_name_by_clo':lambda clo:clo.replace('_1','_0'),
    },
    'viton':{
        'ori_dir':'/mnt/nas/shengjie/datasets/VITON-HD_ori',
        'types':[types[0],],
        # 'types':[types[0],types[0],],
        'clo_dir':['train/cloth'],
        # 'clo_dir':['test/cloth','train/cloth'],
        'human_dir':['train/image'],
        # 'human_dir':['test/image','train/image'],
        'get_human_name_by_clo':lambda clo:clo,
    },
    
}

data_root:
    /mnt/nas/shengjie/datasets/DressCode_1024
clo_type:
    dresses, upper, lower
main_dir:
    $data_root/$clo_type

data_root="/mnt/nas/shengjie/datasets/DressCode_1024"
clo_types=("dresses" "upper" "lower")
sub_dirs=("cloth" "cloth_align_mask-bytedance" "cloth_mask" "densepose" "parse-bytedance" "cloth_align" "cloth_align_parse-bytedance" "cloth_parse" "image" "pose_25")

for clo_type in "${clo_types[@]}"; do
    main_dir="${data_root}/${clo_type}"
    for sub_dir in "${sub_dirs[@]}"; do
        full_sub_dir="${main_dir}/${sub_dir}"
        if [ -d "$full_sub_dir" ]; then
            echo "Processing $full_sub_dir"
            bash restore.sh "$full_sub_dir"
        else
            echo "Directory $full_sub_dir does not exist, skipping."
        fi
    done
done

'''

# INSERT_YOUR_CODE
import argparse
import os

parser = argparse.ArgumentParser()
parser.add_argument('-c', '--cuda', type=str, default=None, help='CUDA_VISIBLE_DEVICES setting')
args, unknown = parser.parse_known_args()

if args.cuda is not None:
    os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda



from PIL import Image
from util_pose import get_pose_predictor
from util_parse import get_parse_predictor
from util_mask_tryon import tryon_pipeline
from tqdm import tqdm

pose_predictor = get_pose_predictor()
parse_predictor = get_parse_predictor()

data_root = "/mnt/nas/shengjie/datasets/DressCode_1024"
clo_types = ["dresses", "upper", "lower"]
sub_dirs = [
    "cloth", "cloth_align_mask-bytedance", "cloth_mask", "densepose", "parse-bytedance",
    "cloth_align", "cloth_align_parse-bytedance", "cloth_parse", "image", "pose_25"
]
get_save_dir = lambda clo_type: f'/mnt/nas/shengjie/datasets/tryon_{clo_type}'

for clo_type in clo_types:
    main_dir = os.path.join(data_root, clo_type)

    human_dir = os.path.join(main_dir, 'image')
    cloth_dir = os.path.join(main_dir, 'cloth')
    cloth_mask_dir = os.path.join(main_dir, 'cloth_mask')

    names_txt = os.path.join(cloth_dir, "names.txt")

    # 读取names.txt，筛选以.jpg结尾的行
    if os.path.isfile(names_txt):
        with open(names_txt, "r") as f:
            jpg_names = [line.strip() for line in f if line.strip().lower().endswith('.jpg')]
        print(f"{names_txt} 中以.jpg结尾的文件名有 {len(jpg_names)} 个")
    else:
        print(f"{names_txt} 不存在，跳过。")

    # INSERT_YOUR_CODE
    # 对于每个clo图片名，找到对应的human图片名，确保两个路径都存在
    for clo_name in tqdm(jpg_names):
        # clo_name: e.g. 000000_1.jpg
        human_name = clo_name.replace('_1', '_0')
        clo_mask_name = clo_name.replace('.jpg', '.png')
        clo_path = os.path.join(cloth_dir, clo_name)
        clo_mask_path = os.path.join(cloth_mask_dir, clo_mask_name)
        human_path = os.path.join(human_dir, human_name)

        if not os.path.isfile(clo_path):
            print(f"Cloth image {clo_path} does not exist, skipping.")
            continue
        if not os.path.isfile(clo_mask_path):
            print(f"Cloth mask image {clo_mask_path} does not exist, skipping.")
            continue
        if not os.path.isfile(human_path):
            print(f"Human image {human_path} does not exist, skipping.")
            continue

        # 可选: 运行tryon_pipeline做一次推理和可视化
        try:
            result = tryon_pipeline(
                img_path=human_path,
                clo_path=clo_path,
                clo_type=clo_type if clo_type!='dresses' else 'full',
                clo_mask_path=clo_mask_path,
                pose_predictor=pose_predictor,
                parse_predictor=parse_predictor
            )
            # 保存可视化结果
            save_dir = get_save_dir(clo_type)
            os.makedirs(save_dir, exist_ok=True)
            out_name = os.path.splitext(clo_name)[0] + "_tryon.jpg"
            out_path = os.path.join(save_dir, out_name)
            result["final_img"].save(out_path)
            print(f"Saved tryon visualization to {out_path}")
            print()
        except Exception as e:
            print(f"Failed tryon for {clo_name}: {e}")

# 后续 还需要在提示词中添加 skin 颜色


'''
cloth RGB 
cloth_align_mask-bytedance  L array([  0, 255], dtype=uint8) cloth's unwarp mask == cloth_mask 
cloth_mask                  L array([  0, 255], dtype=uint8) cloth's unwarp mask
densepose       L       array([ 0,  2,  3,  4,  5,  6,  7,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18,
                                19, 20, 21, 22, 23, 24], dtype=uint8) human's densepose
parse-bytedance
cloth_align RGB == cloth
cloth_align_parse-bytedance  L   array([0, 5], dtype=uint8) cloth's parse
cloth_parse L == cloth_align_parse-bytedance
image RGB human
pose_25 .npy 数据 pose points 可视化为pose图  
        pose 我重新拿模型预测吧
'''